{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We introduce some abilities of OmniGen in this ipynb, including text-to-image, identity-preserving generation, image-conditional generation, and so on.\n",
    "\n",
    "Some tips for generation:\n",
    "- For out of memory or time cost, you can set `offload_model=True` or refer to [./docs/inference.md#requiremented-resources](https://github.com/VectorSpaceLab/OmniGen/blob/main/docs/inference.md#requiremented-resources) to select a appropriate setting.\n",
    "- If the inference time is too long when input multiple images, you can reduce the `max_input_image_size`. More detaild pleae refer to [./docs/inference.md#requiremented-resources](https://github.com/VectorSpaceLab/OmniGen/blob/main/docs/inference.md#requiremented-resources)\n",
    "- Oversaturated: If the image appears oversaturated, please reduce the `guidance_scale`.\n",
    "- Not match the prompt: If the image does not match the prompt, please try to increase the `guidance_scale`.\n",
    "- Low-quality: More detailed prompt will lead to better results. \n",
    "- Animate Style: If the genereate images is in animate style, you can try to add `photo` to the prompt`.\n",
    "- Edit generated image. If you generate a image by omnigen and then want to edit it, you cannot use the same seed to edit this image. For example, use seed=0 to generate image, and should use seed=1 to edit this image.\n",
    "- For image editing tasks, we recommend placing the image before the editing instruction. For example, use `<img><|image_1|></img> remove suit`, rather than `remove suit <img><|image_1|></img>`. \n",
    "- For image editing task and controlnet task, we recommend to set the height and width of output image as the same\n",
    "as input image. For example, if you want to edit a 512x512 image, you should set the height and width of output image as 512x512. You also can set the `use_input_image_size_as_output` to automatically set the height and width of output image as the same as input image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import os\n",
    "# os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # select a gpu to run OmniGen\n",
    "# os.environ['HF_HUB_CACHE'] = 'path_to_save_downloaded_model'\n",
    "\n",
    "from OmniGen import OmniGenPipeline\n",
    "\n",
    "pipe = OmniGenPipeline.from_pretrained(\"Shitao/OmniGen-v1\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 1. Text to Image\n",
    "Here are some examples for text to image generation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Text to Image\n",
    "prompt = [\"A vintage camera placed on the ground, ejecting a swirling cloud of Polaroid-style photographs into the air. The photos, showing landscapes, wildlife, and travel scenes, seem to defy gravity, floating upward in a vortex of motion. The camera emits a glowing, smoky light from within, enhancing the magical, surreal atmosphere. The dark background contrasts with the illuminated photos and camera, creating a dreamlike, nostalgic scene filled with vibrant colors and dynamic movement. Scattered photos are visible on the ground, further contributing to the idea of an explosion of captured memories.\",\n",
    "\"A curly-haired man in a red shirt is drinking tea.\",\n",
    "]\n",
    "for i in range(len(prompt)):\n",
    "    images = pipe(\n",
    "        prompt=prompt[i],  #In fact, you also can pass the entire prompt list here, but it will take more memory cost to generate all images.\n",
    "        height=1024, \n",
    "        width=1024, \n",
    "        guidance_scale=2.5,\n",
    "        separate_cfg_infer=False,\n",
    "        seed=0,\n",
    "        )\n",
    "    # images[0].save(\"i.png\")\n",
    "    images[0].show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 2. Subject-driven Generation or Identity-Preserving Generation\n",
    "\n",
    "You can input an image containing a specific object(eg., human, animal or others), and prompt model to generate a new image based on given object. Different form previous work, OmniGen don't need to detect and crop the object using other models (e.g., segment face is needed in InstandID and PULID). OmniGen will automatically find the specific object in the image and generate image.\n",
    "\n",
    "What'more, OmniGen can process the input images consist of multi objects, our model can automatically identify objects in the image through descriptive instructions, e.g., the right man in `<img><|image_1|></img>`, the woman wearing pink clothes in `<img><|image_2|></img>`.\n",
    "\n",
    "OmniGen also can extract multiple objects from multple images to generate a new image."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "from PIL import Image\n",
    "\n",
    "max_input_image_size = 1024 # you can reduce this size to speed-up the inference and reduce the memory usage\n",
    "\n",
    "prompt=\"The woman in <img><|image_1|></img> waves her hand happily in the crowd\"\n",
    "input_images=[\"./imgs/test_cases/zhang.png\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images, \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.8,  \n",
    "    seed=42\n",
    ")\n",
    "print(\"input_image: \")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "\n",
    "# # You can pass multi objects from multiple input images\n",
    "prompt = \"Two woman are raising fried chicken legs in a bar. A woman is <img><|image_1|></img>. Another woman is <img><|image_2|></img>.\"  # a1\n",
    "input_images = [\"./imgs/test_cases/mckenna.jpg\", \"./imgs/test_cases/Amanda.jpg\"]  \n",
    "images = pipe(\n",
    "        prompt=prompt, \n",
    "        input_images=input_images, \n",
    "        height=1024, \n",
    "        width=1024,\n",
    "        guidance_scale=2.5, \n",
    "        img_guidance_scale=1.8, \n",
    "        max_input_image_size=max_input_image_size,\n",
    "        seed=168\n",
    "    )\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "\n",
    "\n",
    "# # For input images consist of multi objects, our model can automatically identify objects in the image through descriptive instructions. \n",
    "# # e.g., the right man in <img><|image_1|></img>, the woman wearing pink clothes in <img><|image_2|></img>.\n",
    "prompt=\"A man in a black shirt is reading a book. The man is the right man in <img><|image_1|></img>.\"\n",
    "input_images=[\"./imgs/test_cases/two_man.jpg\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.6,\n",
    "    seed=0,\n",
    "    )\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "\n",
    "# # You also can extract multiple objects from multple images to generate a new image\n",
    "# # If you don't describe the clothes in the prompt, the model will tend to retain the clothing from the original image.\n",
    "prompt=\"A man is sitting in the library reading a book, while a woman wearing white shirt next to him is wearing headphone. The man who is reading is the one wearing red sweater in <img><|image_1|></img>. The woman wearing headphones is the right woman wearing suit in <img><|image_2|></img>.\"\n",
    "input_images=[\"./imgs/test_cases/turing.png\", \"./imgs/test_cases/lecun.png\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.8,\n",
    "    max_input_image_size=max_input_image_size,\n",
    "    seed=2)\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "prompt=\"A man and a short-haired woman with a wrinkled face are standing in front of a bookshelf in a library. The man is the man in the middle of <img><|image_1|></img>, and the woman is oldest woman in <img><|image_2|></img>\"\n",
    "input_images=[\"./imgs/test_cases/1.jpg\", \"./imgs/test_cases/2.jpg\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.6,\n",
    "    max_input_image_size=max_input_image_size,\n",
    "    seed=60)\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "\n",
    "prompt=\"A man and a woman are sitting at a classroom desk. The man is the man with yellow hair in <img><|image_1|></img>. The woman is the woman on the left of <img><|image_2|></img>\"\n",
    "input_images=[\"./imgs/test_cases/3.jpg\", \"./imgs/test_cases/4.jpg\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.6,\n",
    "    offload_model=False, # If OOM(out of memory), you can set offload_model=True\n",
    "    max_input_image_size=max_input_image_size,\n",
    "    seed=66)\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "prompt=\"A woman is walking down the street, wearing a white long-sleeve blouse with lace details on the sleeves, paired with a blue pleated skirt. The woman is <img><|image_1|></img>. The long-sleeve blouse and a pleated skirt are <img><|image_2|></img>.\"\n",
    "input_images=[\"./imgs/demo_cases/emma.jpeg\", \"./imgs/demo_cases/dress.jpg\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.6,\n",
    "    offload_model=False, # If OOM(out of memory), you can set offload_model=True\n",
    "    max_input_image_size=max_input_image_size,\n",
    "    seed=666)\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 3. Image-conditional Generation\n",
    "For this task, a representative work is ControlNet. ControlNet requires using other detectors to detect conditions in the input image and then loading the corresponding modules for inference. Unlike ControlNet, OmniGen can complete both tasks (condition extraction and condition-based generation) within a single model and can even achieve this in one step (skipping the condition extraction step)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "# OmniGen can handle some classical CV tasks\n",
    "#prompt = \"Generate the depth map for this image: <img><|image_1|></img>.\"\n",
    "prompt = \"Detect the skeleton of human in this image: <img><|image_1|></img>.\"\n",
    "input_images = [\"./imgs/test_cases/control.jpg\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.6,\n",
    "    separate_cfg_infer=False, # If OOM(out of memory), you can set separate_cfg_infer=True\n",
    "    seed=0)\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "\n",
    "# Generate images based on extracted condition\n",
    "prompt = \"Generate a new photo using the following picture and text as conditions: <img><|image_1|><img>\\n An elderly man wearing gold-framed glasses stands dignified in front of an elegant villa. His gray hair is neatly combed, and his hands rest in the pockets of his dark trousers. He is dressed warmly in a fitted coat over a sweater. The classic villa behind him features ivy-covered walls and large bay windows.\"\n",
    "input_images = [\"./imgs/test_cases/pose.png\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.6,\n",
    "    seed=0)\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()\n",
    "\n",
    "\n",
    "\n",
    "# More simple method to generate images based on image-condition within one step\n",
    "# prompt = \"Following the depth mapping of this image <img><|image_1|></img>, generate a new image: An elderly man wearing a blue hat and gold-framed glasses stands dignified in front of an elegant villa. His gray hair is neatly combed, and his hands rest in the pockets of his dark trousers. He is dressed warmly in a fitted coat over a sweater. The classic villa behind him features ivy-covered walls and large bay windows.\"\n",
    "prompt = \"Following the human pose of this image <img><|image_1|></img>, generate a new photo: An elderly man wearing a gold-framed glasses stands dignified in front of an elegant villa. His gray hair is neatly combed, and his hands rest in the pockets of his dark trousers. He is dressed warmly in a fitted coat over a sweater. The classic villa behind him features ivy-covered walls and large bay windows.\"\n",
    "input_images = [\"./imgs/test_cases/control.jpg\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images,  \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5, \n",
    "    img_guidance_scale=1.6,\n",
    "    seed=0)\n",
    "# images[0].save(\"tii2i.png\")\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 4. Potential Reasoning Capability\n",
    "\n",
    "OmniGen demonstrates a certain level of reasoning capability. It can locate objects based on implicit descriptions rather than explicit instructions, such as identifying something drinkable. This feature could be useful in interesting scenarios, such as robotics."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "\n",
    "prompt = \"<img><|image_1|><\\/img> What item can be used to see the current time? Please remove it.\"\n",
    "input_images = [\"./imgs/test_cases/watch.jpg\"]\n",
    "images = pipe(\n",
    "    prompt=prompt, \n",
    "    input_images=input_images, \n",
    "    height=1024, \n",
    "    width=1024,\n",
    "    guidance_scale=2.5,\n",
    "    img_guidance_scale=1.6,\n",
    "    seed=0)\n",
    "print(f\"input prompt: {prompt}\\ninput image:\")\n",
    "for img in input_images:\n",
    "    Image.open(img).show()\n",
    "print(\"output:\")\n",
    "images[0].show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Other tasks\n",
    "In addition to the above capabilities, OmniGen also offers many other features, such as image editing, denoising, and inpainting. Feel free to explore OmniGen's abilities while using it. \n",
    "\n",
    "If you come up with new tasks, you can also fine-tune OmniGen to give it new capabilities. Fine-tuning is very straightforward, and you can refer to our docs. \n",
    "\n",
    "Have fun!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
